import os
import json
import logging
import requests
import re
from typing import List, Dict, Any, Optional

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("llm_service")

# 添加ollama库导入
try:
    import ollama
except ImportError:
    logger.error("请安装ollama库: pip install ollama")
    raise

# 配置文件路径
CONFIG_FILE = os.path.join(os.path.dirname(__file__), "config.json")

# 默认配置
DEFAULT_CONFIG = {
    "model_name": "deepseek-r1:7b",  # 默认使用deepseek-r1:7b模型
    "max_tokens": 1000,              # 最大生成长度
    "temperature": 0.1,              # 生成多样性
    "show_thinking": True            # 是否显示思考过程
}

# 确保配置文件存在
if not os.path.exists(CONFIG_FILE):
    with open(CONFIG_FILE, "w", encoding="utf-8") as f:
        json.dump(DEFAULT_CONFIG, f, ensure_ascii=False, indent=2)

def load_config():
    """加载LLM配置"""
    if os.path.exists(CONFIG_FILE):
        with open(CONFIG_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    return DEFAULT_CONFIG

def save_config(config):
    """保存LLM配置"""
    with open(CONFIG_FILE, "w", encoding="utf-8") as f:
        json.dump(config, f, ensure_ascii=False, indent=2)

class LLMService:
    """大模型服务类"""
    
    def __init__(self):
        self.config = load_config()
        # 启动时检查模型可用性
        self._check_model_availability()
    
    def _check_model_availability(self):
        """检查配置的模型是否可用，如果不可用则使用默认模型"""
        try:
            logger.info("使用默认模型：deepseek-r1:7b")
            self.config["model_name"] = "deepseek-r1:7b"
            save_config(self.config)
        except Exception as e:
            logger.error(f"检查模型可用性时出错: {str(e)}")
    
    def update_config(self, new_config):
        """更新配置"""
        self.config.update(new_config)
        save_config(self.config)
    
    def get_config(self):
        """获取当前配置"""
        return self.config
    
    def get_available_models(self):
        """获取本地可用的模型列表"""
        # 直接返回默认模型，因为本地已经确认运行了该模型
        logger.info("直接使用本地已有的deepseek-r1:7b模型")
        return ["deepseek-r1:7b"]
    
    def generate_response(self, prompt: str, knowledge_texts: List[str] = None) -> str:
        """
        根据提示和知识生成响应
        
        Args:
            prompt: 用户输入的问题或提示
            knowledge_texts: 检索到的相关知识条目列表
            
        Returns:
            大模型生成的回答
        """
        # 构建系统提示
        system_prompt = """你是一个基于知识库的智能助手。请根据所提供的知识信息回答问题。
        
如果知识库中没有相关信息，请明确指出。
这里是我给用户的最终回答...
"""
        
        # 打印系统提示
        logger.info(f"系统提示: {system_prompt}")
        
        # 构建带有知识的用户提示
        if knowledge_texts and len(knowledge_texts) > 0:
            # 打印检索到的知识
            logger.info("检索到的知识:")
            for i, text in enumerate(knowledge_texts):
                logger.info(f"  知识{i+1}: {text}")
                
            knowledge_str = "\n".join([f"- {text}" for text in knowledge_texts])
            user_prompt = f"""
请根据以下知识回答问题:

{knowledge_str}

问题: {prompt}

请先在<think>标签中分析上述知识，推理出合理的结论。然后给出最终回答。
回答格式请参考系统提示。
"""
        else:
            user_prompt = f"""
问题: {prompt}

请先在<think>标签中分析这个问题并进行推理。然后给出最终回答。
回答格式请参考系统提示。
"""
            logger.info("未检索到相关知识，直接使用用户问题作为提示")
        
        # 打印完整的用户提示(截断过长内容)
        if len(user_prompt) > 200:
            logger.info(f"用户提示(截断): {user_prompt[:200]}...")
        else:
            logger.info(f"用户提示: {user_prompt}")
        
        try:
            # 获取要使用的模型
            model_name = self.config.get("model_name", "deepseek-r1:7b")
            logger.info(f"使用模型 '{model_name}' 生成回答")
            
            # 调用模型
            response = ollama.chat(
                model=model_name,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                options={
                    "temperature": self.config.get("temperature", 0.1),
                    "num_predict": self.config.get("max_tokens", 1000)
                }
            )
            
            # 提取生成的回复
            if response and "message" in response and "content" in response["message"]:
                content = response["message"]["content"]
                
                # 如果不显示思考过程，完全移除思考部分及其标记
                if not self.config.get("show_thinking", True) and "<thinking>" in content:
                    content = re.sub(r'<think>[\s\S]*?</think>', '', content).strip()
                
                return content
            else:
                return "模型未返回有效响应"
        
        except Exception as e:
            error_msg = str(e)
            logger.error(f"调用模型时出错: {error_msg}")
            
            # 如果是模型未找到错误，尝试使用第一个可用模型
            if "not found" in error_msg.lower() and "404" in error_msg:
                return self._retry_with_available_model(system_prompt, user_prompt)
            
            return f"调用本地模型时出错: {error_msg}"
    
    def _retry_with_available_model(self, system_prompt, user_prompt):
        """尝试使用可用模型重试请求"""
        try:
            # 获取可用模型
            models = self.get_available_models()
            if not models or models[0].startswith("获取模型列表出错"):
                return "无法获取可用模型列表"
            
            model_name = models[0]
            logger.info(f"尝试使用可用模型: {model_name}")
            
            # 更新配置并保存
            self.config["model_name"] = model_name
            save_config(self.config)
            
            # 重新尝试
            response = ollama.chat(
                model=model_name,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                options={
                    "temperature": self.config.get("temperature", 0.1),
                    "num_predict": self.config.get("max_tokens", 1000)
                }
            )
            
            if response and "message" in response and "content" in response["message"]:
                content = response["message"]["content"]
                
                # 如果不显示思考过程，完全移除思考部分及其标记
                if not self.config.get("show_thinking", True) and "<think>" in content:
                    content = re.sub(r'<think>[\s\S]*?</think>', '', content).strip()
                
                return content
            else:
                return "模型未返回有效响应"
        
        except Exception as retry_error:
            logger.error(f"使用备用模型重试失败: {str(retry_error)}")
            return f"使用备用模型失败: {str(retry_error)}" 